#!/usr/bin/env python
# encoding:utf-8

"""
@author: DeltaF
@software: pycharm
@file: ma_strategy.py
@time: 2021/9/12 15:54
@desc:双均线策略
"""

import datetime
import data.stock as st
import strategy.base as strat
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt



def ma_strategy(data, short_window=5, long_window=20):
    """
    多均线策略
    :param data:
    :param short_window:
    :param long_window:
    :return:
    """

    data = pd.DataFrame(data)
    # 计算技术指标：ma短期，ma长期
    data['short_ma'] = data['close'].rolling(window=short_window).mean()
    data['long_ma'] = data['close'].rolling(window=long_window).mean()

    # 生成信号：金叉买入，死叉卖出
    data['buy_signal'] = np.where(data['short_ma'] > data['long_ma'], 1, 0)
    data['sell_signal'] = np.where(data['short_ma'] < data['long_ma'], -1, 0)
    print(data[['close', 'short_ma', 'long_ma', 'buy_signal', 'sell_signal']])

    # 过滤信号：st.compose_signal
    data = strat.compose_signal(data)
    #print(data[['close', 'short_ma', 'long_ma', 'signal']])

    # 计算单次收益
    data = strat.calculate_prof_pct(data)
    #print(data)

    # 计算累计收益率
    data = strat.calculate_cum_prof(data)

    # 数据预览
    # print(data[['close', 'short_ma', 'long_ma', 'signal', 'cum_profit']])

    return data


if __name__ == '__main__':
    # 股票列表
    stock = ['000001.XSHE', '000858.XSHE', '002594.XSHE']
    df = st.get_single_price('601888.XSHG', 'daily', '2021-01-01', datetime.datetime.now())
    df = ma_strategy(df)  # 调用双均线策略
    # for code in stocks:
    #     df = st.get_single_price(code, 'daily', '2016-01-01', '2021-01-01')
    #     df = ma_strategy(df)  # 调用双均线策略
